`timescale 1ns / 1ps
//////////////////////////////////////////////////////////////////////////////////
// Company: 
// Engineer: 
// 
// Create Date: 2025/10/02 11:56:56
// Design Name: 
// Module Name: fc_1
// Project Name: 
// Target Devices: 
// Tool Versions: 
// Description: 
// 
// Dependencies: 
// 
// Revision:
// Revision 0.01 - File Created
// Additional Comments:
// 
//////////////////////////////////////////////////////////////////////////////////


module fc_1 (
    input  wire         clk_i,
    input  wire         rst_n_i,
    input  wire         start_i,
    output wire         done_o,

    output reg  [1:0]    input_data_addr_o,
    input  wire [1023:0] input_data_i,          // 128 * 8-bit

    output reg  [9:0]    weight_addr_o,         // {out_channel[6:0], group_idx[1:0]}
    input  wire [1023:0] weight_i,              // 128 * 8-bit

    output reg  [6:0]    bias_addr_o,
    input  wire [7:0]    bias_i,                // 8-bit signed

    output reg           fc_output_wren_o,
    output reg  [31:0]   fc_output_data_o,
    output reg  [6:0]    fc_output_addr_o,

    output wire [1023:0] mul_data1_o,
    output wire [1023:0] mul_data2_o,
    input  wire [2047:0] mul_result_i           // 128 * 16-bit (8x8)
);
    parameter INPUT_DIM    = 512;
    parameter OUTPUT_DIM   = 128;
    parameter GROUP_SIZE   = 128;
    parameter GROUPS       = 4;
    parameter MULT_LATENCY = 3;
    parameter RAM_LATENCY  = 3;

    // 主循环计数
    reg [6:0] out_channel;
    reg [1:0] group_idx;
    reg       running;

    // 控制信号延迟链：使用参数化索引，覆盖到 DONE 阶段
    // 实际所需最大索引 = INDEX_DONE + 2 (冗余再给几拍)
    reg [6:0] out_channel_dly [0:24];
    reg [1:0] group_idx_dly   [0:24];

    // 偏置同样来自 BRAM (3 拍延迟)，后续再加移位链对齐
    // 估计最大需要的偏置移位级数 = (INDEX_BIAS - RAM_LATENCY)
    reg [7:0]    bias_reg  [0:24];

    // 加法树结果    
    wire [31:0] final_sum;

    // ================== 流水线阶段索引参数化 ==================
    // Address 发起为 0
    localparam INDEX_ADDR      = 0;                                  // 地址阶段
    localparam INDEX_RAM_OUT   = INDEX_ADDR + RAM_LATENCY;           // RAM 输出可用于乘法输入 (3)
    localparam INDEX_MUL_OUT   = INDEX_RAM_OUT + MULT_LATENCY;       // 乘法器输出 (3+3=6)
    localparam ADDER_TREE_STAGES = 3;                                // adder_tree 内部寄存级数
    localparam INDEX_ADDER_OUT = INDEX_MUL_OUT + ADDER_TREE_STAGES;  // 加法树输出 (6+3=9)
    localparam INDEX_ACCUM     = INDEX_ADDER_OUT;                    // 分组累加
    localparam INDEX_BIAS      = INDEX_ACCUM + 1;                    // 加偏置
    localparam INDEX_RELU      = INDEX_BIAS + 1;                     // ReLU
    localparam INDEX_WRITE     = INDEX_RELU + 1;                     // 写回
    localparam INDEX_DONE      = INDEX_WRITE + 1;                    // done 输出

    // 偏置有效索引（其对应地址在 INDEX_ADDR 发起，BRAM 3 拍输出）
    localparam BIAS_PIPE_OFFSET = (INDEX_BIAS > RAM_LATENCY) ? (INDEX_BIAS - RAM_LATENCY) : 0;

    // 每输出通道跨4组累加
    reg  [31:0] channel_accum [0:OUTPUT_DIM-1];

    // 偏置后累加器 / ReLU
    reg [31:0] accumulator;
    reg [31:0] relu_result;

    // done 产生移位寄存器 (last_launch 后延迟到 INDEX_DONE)
    reg [INDEX_DONE:0] done_shift; // 位宽自适应

    integer i;

    // 启动 / 运行控制
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            running <= 1'b0;
        end else begin
            if(start_i && !running)
                running <= 1'b1;
            else if(done_o) // 流水线完全结束
                running <= 1'b0;
        end
    end

    // group_idx 更新
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            group_idx <= 0;
        end else if(running) begin
            if(group_idx == GROUPS-1)
                group_idx <= 0;
            else
                group_idx <= group_idx + 1'b1;
        end
    end

    // out_channel 更新：在 group 回卷处前进
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            out_channel <= 0;
        end else if(running && group_idx == GROUPS-1) begin
            if(out_channel == OUTPUT_DIM-1)
                out_channel <= 0; // 结束后无所谓 (下次 start 会复用)
            else
                out_channel <= out_channel + 1'b1;
        end
    end

    // 延迟链 (简单移位)
    integer d;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            for(d=0; d<=24; d=d+1) begin
                out_channel_dly[d] <= 0;
                group_idx_dly[d]   <= 0;
            end
        end else begin
            out_channel_dly[0] <= out_channel;
            group_idx_dly[0]   <= group_idx;
            for(d=1; d<=24; d=d+1) begin
                out_channel_dly[d] <= out_channel_dly[d-1];
                group_idx_dly[d]   <= group_idx_dly[d-1];
            end
        end
    end

    // 阶段1：地址生成
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            input_data_addr_o <= 0;
            weight_addr_o     <= 0;
            bias_addr_o       <= 0;
        end else if(running) begin
            input_data_addr_o <= group_idx;
            weight_addr_o     <= {out_channel, group_idx}; // 7+2 -> 9 bits (上位补0给10位)
            bias_addr_o       <= out_channel;
        end
    end

    // 阶段2-4：RAM 输出寄存对齐 (3拍)
    
    // 偏置延迟链：bias_i 已包含 BRAM 3 拍延迟，这里再移位用于对齐到 INDEX_BIAS
    integer b;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            for(b=0;b<=24;b=b+1) bias_reg[b] <= 0;
        end else begin
            bias_reg[0] <= bias_i;
            for(b=1;b<=24;b=b+1)
                bias_reg[b] <= bias_reg[b-1];
        end
    end

    // 乘法器输入
    assign mul_data1_o = input_data_i;
    assign mul_data2_o = weight_i;

    // 阶段5-7：乘法结果延迟

    // 阶段8-10：加法树 (3级流水) 
    adder_tree_1 #(
        .NUM_INPUTS(128),
        .DATA_WIDTH(16),
        .OUT_WIDTH(32)
    ) u_adder_tree (
        .clk    (clk_i),
        .rst_n  (rst_n_i),
        .data   (mul_result_i),
        .sum_out(final_sum)
    );

    // 阶段11：分组累加：在 INDEX_ACCUM 周期 final_sum 对应 group_idx_dly[INDEX_ACCUM]
    integer ca;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            for(ca=0; ca<OUTPUT_DIM; ca=ca+1)
                channel_accum[ca] <= 0;
        end else begin
            // FIXME 数学上好像可以不用区分第一组？
            if(group_idx_dly[INDEX_ACCUM] == 0) begin
                // 第一组：写入
                channel_accum[out_channel_dly[INDEX_ACCUM]] <= final_sum;
            end else begin
                // 后续组：累加
                channel_accum[out_channel_dly[INDEX_ACCUM]] <= channel_accum[out_channel_dly[INDEX_ACCUM]] + final_sum;
            end
        end
    end

    // 阶段12：加偏置 (INDEX_BIAS) 仅在最后一组
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            accumulator <= 0;
        end else begin
            if(group_idx_dly[INDEX_BIAS] == (GROUPS-1)) begin
                accumulator <= channel_accum[out_channel_dly[INDEX_BIAS]] +
                               {{24{bias_reg[BIAS_PIPE_OFFSET][7]}}, bias_reg[BIAS_PIPE_OFFSET]};
            end else begin
                accumulator <= 0;
            end
        end
    end

    // 阶段13：ReLU (INDEX_RELU)
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i)
            relu_result <= 0;
        else if(accumulator[31])
            relu_result <= 0;
        else
            relu_result <= accumulator;
    end

    // 阶段14：写输出 (INDEX_WRITE)
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            fc_output_wren_o <= 0;
            fc_output_data_o <= 0;
            fc_output_addr_o <= 0;
        end else begin
            if(group_idx_dly[INDEX_WRITE] == (GROUPS-1)) begin
                fc_output_wren_o <= 1'b1;
                fc_output_addr_o <= out_channel_dly[INDEX_WRITE];
                fc_output_data_o <= relu_result;
            end else begin
                fc_output_wren_o <= 1'b0;
            end
        end
    end

    // FIXME 为什么要写这么复杂
    // 阶段15：done 产生：last_launch -> 在 INDEX_DONE 周期输出
    wire last_launch = running &&
                       (out_channel == OUTPUT_DIM-1) &&
                       (group_idx == GROUPS-1);
    integer ds;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            done_shift <= 0;
        end else begin
            done_shift <= {done_shift[INDEX_DONE-1:0], last_launch};
        end
    end
    assign done_o = done_shift[INDEX_DONE];
endmodule
